import paramiko
import sys


class SSHTunnel:
    def __init__(self, ssh_host, ssh_port, ssh_user, ssh_password, remote_host, remote_port, local_port):
        self.ssh_host = ssh_host
        self.ssh_port = ssh_port
        self.ssh_user = ssh_user
        self.ssh_password = ssh_password
        self.remote_host = remote_host
        self.remote_port = remote_port
        self.local_port = local_port
        self.tunnel = None

    def create_tunnel(self):
        # 创建SSH对象
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            ssh.connect(self.ssh_host, username=self.ssh_user, password=self.ssh_password, port=self.ssh_port)
        except paramiko.SSHException as e:
            print(f"建立SSH连接失败: {e}")
            return

        # 建立隧道
        try:
            self.tunnel = paramiko.Transport((self.ssh_host, self.ssh_port))
            self.tunnel.connect(username=self.ssh_user, password=self.ssh_password)
            local_addr = ('127.0.0.1', self.local_port)
            remote_addr = (self.remote_host, self.remote_port)
            forward_tunnel = paramiko.SFTPClient._open_forward_agent_channel(self.tunnel, remote_addr, local_addr)
            print("SSH隧道已建立")
        except paramiko.SSHException as e:
            print(f"建立SSH隧道失败: {e}")
            return

    def close_tunnel(self):
        if self.tunnel:
            self.tunnel.close()
            print("SSH隧道已关闭")


#用于进行测试
if __name__ == "__main__":
    # SSH服务器的配置
    ssh_host = 'your_ssh_host'
    ssh_port = 22
    ssh_user = 'your_ssh_user'
    ssh_password = 'your_ssh_password'

    # 远端服务器的配置
    remote_host = 'remote_host_to_tunnel'
    remote_port = 1234  # 远端服务器的端口
    local_port = 8080  # 本地映射的端口

    # 创建SSH隧道
    tunnel = SSHTunnel(ssh_host, ssh_port, ssh_user, ssh_password, remote_host, remote_port, local_port)
    try:
        tunnel.create_tunnel()
        # 隧道建立后，你可以在这里执行其他操作
        # 例如，你可以使用requests库来请求localhost的local_port
        import requests
        response = requests.get(f'http://127.0.0.1:{local_port}/users/class/getClasses')
        print(response.text)
    except Exception as e:
        print(f"SSH隧道创建失败: {e}")
    finally:
        tunnel.close_tunnel()